import scapy.all as scapy
import os
import scapy_http.http

wirshark_data = []


def parse_pcap():  #提取五元组
    path = "D:\python_Data\Brian_IDS\static\pcap"
    files = os.listdir(path)
    for file in files:  # 遍历文件夹
        packets = scapy.rdpcap(path + '/' + file)
        for p in packets:
            try:
                wirshark_data.append(str(p[1].dst)+' '+str(p[1].dport)+' '+str(p[1].src)+' '+str(p[1].sport)+' '+str(p[1].proto)+' '+str(p[1].Host.decode("utf-8")))
            except AttributeError:
                try:
                    wirshark_data.append(
                        str(p[1].dst) + ' ' + str(p[1].dport) + ' ' + str(p[1].src) + ' ' + str(p[1].sport) + ' ' + str(p[1].proto))
                except AttributeError:
                    continue
def alert(threat_data,cheak_data):
    f = open("D:\python_Data\Brian_IDS\static/rules\ids-ioc-result.txt", "a")
    if threat_data.split(' ')[4]:
        threat_data.split(' ')[4] ='tcp'

    try:                                #域名
        destination = threat_data.split(' ')[0]
        dport = threat_data.split(' ')[1]
        source = threat_data.split(' ')[2]
        sport = threat_data.split(' ')[3]
        proto = threat_data.split(' ')[4]=='6' and 'tcp' or 'udp'
        dns = threat_data.split(' ')[5]
        print('[+]检测到危险域名：'+dns+'')
        print("[*]目的地址："+destination+'\t目的端口：'+dport+'\t源地址：'+source+'\t源端口：'+sport+'\t协议：'+proto+ '\n')
        f.write('[+]检测到危险域名：' + dns + '\n')
        f.write("[*]目的地址："+destination+'\t目的端口：'+dport+'\t源地址：'+source+'\t源端口：'+sport+'\t协议：'+proto)
        f.write('\n')
    except:             #IP+sport
        destination = threat_data.split(' ')[0]
        dport = threat_data.split(' ')[1]
        source = threat_data.split(' ')[2]
        sport = threat_data.split(' ')[3]
        proto = threat_data.split(' ')[4]=='6' and 'tcp' or 'udp'
        print('[+]检测到危险IP及相应端口：'+cheak_data)
        print("[*]目的地址：" + destination + '\t目的端口：' + dport + '\t源地址：' + source + '\t源端口：' + sport + '\t协议：' + proto+ '\n')
        f.write('[+]检测到危险IP及相应端口：'+cheak_data+'\n')
        f.write("[*]目的地址：" + destination + '\t目的端口：' + dport + '\t源地址：' + source + '\t源端口：' + sport + '\t协议：' + proto)
        f.write('\n')
    f.close()


def check():
    with open('D:\python_Data\Brian_IDS\static/rules\Threat_Intelligence_base.txt','r') as file:
        pcap_Check = file.readlines()
        for i in range(0,len(wirshark_data)-1):
            for pcap_one in pcap_Check:
                if pcap_one.split('\n')[0] in wirshark_data[i]:
                    alert(str(wirshark_data[i]),pcap_one.split('\n')[0])

if __name__ == '__main__':
    parse_pcap()
    check()
    print("[*]检测完毕，数据已存储到”D:\python_Data\Brian_IDS\static/rules\ids-ioc-result.txt“")
